"""
脚本名称：download_models.py
用途：使用 modelscope 批量下载多个大模型

✅ 使用说明：

1. 安装依赖：
   pip install modelscope

2. 执行脚本：
   python download_models.py --models "Qwen/Qwen2.5-Omni-7B"

3. 下载位置说明：
   所有模型将被下载到 cache_dir 指定的本地目录中（例如：/data/AE/llm/models/）

4. 下载完成后，终端会输出如下信息：
   模型 {model_id} 已下载至目录: {model_dir}
   模型 {model_id} 下载耗时: {model_download_time:.2f}秒

备注：建议确保本地磁盘空间充足，模型体积较大。

"""

import time
import argparse
from modelscope import snapshot_download

parser = argparse.ArgumentParser(description='批量下载模型')
parser.add_argument('--models', type=str, required=True, help='逗号分隔的模型ID列表')
parser.add_argument('--cache_dir', type=str, default='/data/AE/llm/models/', help='模型缓存目录')

args = parser.parse_args()
model_ids = args.models.split(',')
cache_dir = args.cache_dir
start_time = time.time()

# 遍历模型列表，依次下载
for model_id in model_ids:
    model_start_time = time.time()
    model_dir = snapshot_download(model_id, cache_dir=cache_dir)
    model_end_time = time.time()
    model_download_time = model_end_time - model_start_time
    print(f'模型 {model_id} 已下载至目录: {model_dir}')
    print(f'模型 {model_id} 下载耗时: {model_download_time:.2f}秒')

end_time = time.time()
total_time = end_time - start_time
print(f'\n所有模型下载完成，总耗时: {total_time:.2f}秒')
